# 数据中心：从data文件夹中提取数据
import datetime
import os
import numpy as np
import pandas as pd
from pandas import read_parquet
import warnings

from algo_features_processor import FTProcessor

warnings.filterwarnings('ignore')


class DataCenter:
    date_init = datetime.date(2018, 12, 31)
    date_end = datetime.date(2024, 12, 31)

    def __init__(self):
        if os.path.exists(f"data/cleaned_universe.parquet"):
            self.uni = read_parquet(f"data/cleaned_universe.parquet")
            self.fin_field = self.uni.columns[-38:].tolist()
            self.dates = self.uni['date'].unique().tolist()
        else:
            self.generated_from_raw()

    def generated_from_raw(self):
        # 股票池数据
        self.uni = read_parquet(f"data/中证800成分股权重.parquet").drop(columns=['指数代码']).rename(columns={'股票代码': 'code', '交易日期': 'date', '权重': 'weight'})
        self.uni = self.uni[self.uni['date'] >= self.date_init]
        self.uni = self.uni[self.uni['date'] <= self.date_end]

        # 行情数据
        self.mkt: pd.DataFrame = pd.read_feather(f"data/A股行情数据.feather").rename(columns={
            '股票代码': 'code',
            '交易日期': 'date',
            '开盘价': 'open',
            '最高价': 'high',
            '最低价': 'low',
            '收盘价': 'close',
            'VWAP': 'vwap',
            '成交量': 'volume',
            '成交额': 'amount',
            '复权因子': 'adj_factor',
            '交易状态': 'is_trade',
        })
        self.mkt.sort_values(['date', 'code'], inplace=True)
        self.mkt = self.mkt[self.mkt['date'] >= (self.date_init - datetime.timedelta(days=31))]

        # 日期
        self.dates: np.ndarray = self.mkt['date'].unique()

        # 处理权重->股票池
        self.uni = self.uni.groupby('date', as_index=False).apply(lambda x: self._fill_weights(x, self.dates)).reset_index(drop=True)
        self.uni = self.uni.merge(self.mkt, how='left', on=['date', 'code'])
        self.uni = self.uni[self.uni['is_trade'] == '交易']
        self.uni.drop(columns=['is_trade'], inplace=True)
        del self.mkt

        # 市值数据
        self.size = pd.read_feather(f"data/A股总市值数据.feather").rename(columns={'股票代码': 'code', '交易日期': 'date', '总市值':'tot_cap'})
        self.size = FTProcessor.fillna_ts_expand(self.size)
        self.size = self.size[self.size['date'] >= self.date_init]
        self.size = self.size[self.size['date'] <= self.date_end]
        self.uni = self.uni.merge(self.size, how='left', on=['date', 'code'])
        del self.size

        # 行业数据
        self.indus = read_parquet(f"data/A股中信一级行业分类.parquet").stack().to_frame(name='sector').reset_index().rename(
            columns={'s_info_windcode': 'code', 'trade_dt': 'date'})
        self.indus = self.indus[self.indus['date'] >= self.date_init]
        self.indus = self.indus[self.indus['date'] <= self.date_end]
        self.uni = self.uni.merge(self.indus, how='left', on=['date', 'code'])
        del self.indus

        # 基本面eod
        self.fin: pd.DataFrame = read_parquet(f"data/基本面因子.parquet").reset_index().rename(columns={'s_info_windcode': 'code', 'trade_dt': 'date'})
        self.fin = self.fin[self.fin['date'] >= self.date_init]
        self.fin = self.fin[self.fin['date'] <= self.date_end]
        for k in self.fin.columns[2:]:
            self.fin.loc[:, k] = FTProcessor.fillna_ts_expand(self.fin[['date', 'code', k]])[k].values
        self.uni = self.uni.merge(self.fin, how='left', on=['date', 'code'])
        self.fin_field = self.fin.columns[2:].tolist()
        del self.fin

        self.uni.to_parquet(f"data/cleaned_universe.parquet", engine='pyarrow', index=False)

    def read_idx_eod(self):
        # 指数eod
        self.idx: pd.DataFrame = read_parquet(f"data/中证800收盘价.parquet")
        self.idx.columns.name = None
        self.idx.index.name = 'date'
        self.idx = self.idx[self.idx.index >= self.date_init]
        self.idx = self.idx[self.idx.index <= self.date_end]
        self.idx = self.idx.pct_change().dropna().sort_values('date')

    @staticmethod  # 填充月间权重
    def _fill_weights(x: pd.DataFrame, dates: np.ndarray) -> pd.DataFrame:
        init_date = x['date'].min()
        month = init_date.month+2
        year = init_date.year
        if month > 12:
            month = month%12
            year = year+1
        end_date = min(datetime.date(year, month ,1) - datetime.timedelta(days=1), datetime.date(2024, 11,29))
        curr_dates = dates[dates > init_date]
        curr_dates = curr_dates[curr_dates <= end_date]
        if len(curr_dates) > 0:
            df = []
            for i in curr_dates:
                x_copy = x.copy(deep=False)
                x_copy['date'] = i
                df.append(x_copy)
            df = pd.concat(df)
            return df
        else:
            return pd.DataFrame()


if __name__ == '__main__':
    dc = DataCenter()
    print(dc.uni)
